package io.sqrtqiezi.spark.dataframe

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions._

object AggregateOperator {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession
      .builder()
      .master("local")
      .appName("aggregate practise")
      .getOrCreate()

    val df = spark.read.format("csv")
      .option("header", "true")
      .option("inferSchema", "true")
      .load("data/retail-data/all/*.csv")
      .coalesce(5)

    df.cache
    df.createOrReplaceTempView("dfTable")

    df.show
    println(df.count)

    df.select(countDistinct("StockCode")).show

    df.selectExpr("count(distinct StockCode) as count").show

    df.select(approx_count_distinct("StockCode", 0.1)
      .alias("count_stockcode")).show

    df.select(first("StockCode"),
      last("StockCode"),
      max("Quantity"),
      min("Quantity")).show

    df.select(
      count("Quantity").alias("total_transactions"),
      sum("Quantity").alias("total_purchases"),
      avg("Quantity").alias("avg_purchases"),
      expr("mean(Quantity)").alias("mean_purchases")
    )
      .selectExpr(
        "total_purchases/total_transactions",
        "avg_purchases",
        "mean_purchases"
      )
      .show

    df.agg(collect_set("Country"), collect_list("Country"))
      .show

    df.groupBy("InvoiceNo", "CustomerId").count.orderBy("count").show

    // 窗口函数
    val dfWithDate = df.withColumn("date", to_date(col("InvoiceDate"), "MM/d/yyyy H:mm"))

    val windowSpec = Window.partitionBy("CustomerId", "date")
      .orderBy(col("Quantity").desc)
      .rowsBetween(Window.unboundedPreceding, Window.currentRow)

    val maxPurchaseQuantity = max(col("Quantity")).over(windowSpec)
    val purchaseDenseRank = dense_rank().over(windowSpec)
    val purchaseRank = rank().over(windowSpec)

    println(purchaseDenseRank.explain(true))

    val df3 = dfWithDate.where("CustomerId is not null")
      .orderBy(col("CustomerId"))
      .select(
        col("CustomerId"),
        col("date"),
        col("Quantity"),
        purchaseRank.alias("quantityRank"),
        purchaseDenseRank.alias("quantityDenseRank"),
        maxPurchaseQuantity.alias("maxPurchaseQuantity")
      )

    println(df3.explain)

    df3.show

    // 分组集
    val dfNotNull = dfWithDate.drop()
    val rolledUpDF = dfNotNull.rollup("Date", "Country")
      .agg(sum("Quantity"))
      .selectExpr("Date", "Country", "`sum(Quantity)` as total_quantity")
      .orderBy("Date")

    rolledUpDF.show

    rolledUpDF.where("Country is null").show
    rolledUpDF.where("Date is null").show
  }
}
